Block Sparse Attention with Log-Linear Complexity
金字塔式稀疏注意力：对数线性复杂度
🧭 长上下文 LLM 稀疏注意力加速 ★ 深度解读
Shanghai Jiao Tong University · Shanghai Innovation Institute · ByteDance Seed
PISA 用金字塔 Top-K 块选择 + LogSumExp 评分，把可训练稀疏注意力的选择复杂度从 O(N^2) 降至 O(N log N)
📌 工作参考价值（4 点）
- 用户行为序列、商品文档、Agent 长记忆等长上下文场景，可以将 block sparse attention 的 Top-K 选择换成 PISA 的层级候选缩小：长序列（>64K）下路由 latency 显著低于普通 BSA，适合海量长行为序列或长文档理解；短序列（<16K）仍用普通 BSA 更划算。
- 工程实现可借鉴两段式 kernel：训练/prefill 阶段先做 intermediate-level 路由，再把同一 leaf key block 的 query 聚成 tile 复用 key IO，Qtile=4 时 64K+ 额外获得约 1.3× 加速；解码阶段用单 kernel 融合所有层级避免多次 launch，缓存 mean pyramid 只更新当前 token 的祖先路径，均摊 O(log N)。
- 块打分用 LogSumExp 而不用均值或均值+方差：LSE 保留更高阶信息，提升 Top-K 块选择质量和检索类/containment 任务准确率；在 GQA 下按 query heads 求和 LSE 做共享块选择可复用，但要注意它和 full attention 的 head-mass 平均排序规则不完全一致。
- 若业务中有 Agent/RAG 长候选集粗筛，可把“池化层级 + 每层 bounded Top-K 扩展”作为两阶段检索范式：先粗粒度过滤，再对少数候选精确打分，避免全量候选打分；训练时梯度只经过选中 attention，离散选择不参与反传，可作为稀疏模块插入现有结构。